# main.py
import asyncio
from config import Config
from db import DB
from ai_client import AIClientAsync
from name_generator import NameGenerator
from name_completer import NameCompleter
import argparse

async def main():
    parser = argparse.ArgumentParser(description="AI Name Generator")
    parser.add_argument("action", choices=["generate", "complete"], help="Action to perform: 'generate' or 'complete'")

    # generate 命令的参数
    parser.add_argument("--reset", action="store_true", help="Reset all surnames' is_active to True (only for 'generate')")
    parser.add_argument("--surname", type=str, help="Specific surname to generate names for (only for 'generate')")
    parser.add_argument("--length", type=int, choices=[2, 3], help="Specific name length (2 or 3) (only for 'generate')")
    parser.add_argument("--style", type=str, help="Specific style to generate names (only for 'generate')")

    # complete 命令的参数
    parser.add_argument("--complete-surname", type=str, dest="complete_surname", help="Specific surname to complete details for (only for 'complete')")
    parser.add_argument("--complete-length", type=int, choices=[2, 3], dest="complete_length", help="Specific name length (2 or 3) to complete details for (only for 'complete')")
    parser.add_argument("--complete-style", type=str, dest="complete_style", help="Specific style to complete details for (only for 'complete')")
    parser.add_argument("--start-id", type=int, dest="start_id", help="Start ID for completing details (only for 'complete')") # 新增参数


    args = parser.parse_args()

    config = Config()
    db = DB(config)


    if args.action == "generate":
        # 为 generate 创建 AIClientAsync，并指定 task_type
        ai_client_generate = AIClientAsync(config, task_type="generate")
        name_generator = NameGenerator(config, db, ai_client_generate)
        await name_generator.generate_names(reset=args.reset, surname=args.surname, length=args.length, style=args.style)

    elif args.action == "complete":
        # 为 complete 创建 AIClientAsync，并指定 task_type
        ai_client_complete = AIClientAsync(config, task_type="complete")
        name_completer = NameCompleter(config, db, ai_client_complete)
        await name_completer.complete_details(surname=args.complete_surname, length=args.complete_length, style=args.complete_style, start_id=args.start_id) # 传递 start_id
    else:
        print("Invalid action.")

    await db.close()

if __name__ == "__main__":
    asyncio.run(main())